{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "matchzoo version 1.0\n",
      "`ranking_task` initialized with metrics [normalized_discounted_cumulative_gain@3(0.0), normalized_discounted_cumulative_gain@5(0.0), mean_average_precision(0.0)]\n",
      "data loading ...\n",
      "data loaded as `train_pack_raw` `dev_pack_raw` `test_pack_raw`\n"
     ]
    }
   ],
   "source": [
    "%run init.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "ranking_task = mz.tasks.Ranking(losses=mz.losses.RankCrossEntropyLoss(num_neg=1))\n",
    "ranking_task.metrics = [\n",
    "    mz.metrics.NormalizedDiscountedCumulativeGain(k=3),\n",
    "    mz.metrics.NormalizedDiscountedCumulativeGain(k=5),\n",
    "    mz.metrics.MeanAveragePrecision()\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "preprocessor = mz.models.aNMM.get_default_preprocessor()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing text_left with chain_transform of Tokenize => Lowercase => PuncRemoval: 100%|██████████| 2118/2118 [00:00<00:00, 9296.68it/s]\n",
      "Processing text_right with chain_transform of Tokenize => Lowercase => PuncRemoval: 100%|██████████| 18841/18841 [00:03<00:00, 5986.25it/s]\n",
      "Processing text_right with append: 100%|██████████| 18841/18841 [00:00<00:00, 542075.71it/s]\n",
      "Building FrequencyFilter from a datapack.: 100%|██████████| 18841/18841 [00:00<00:00, 143392.21it/s]\n",
      "Processing text_right with transform: 100%|██████████| 18841/18841 [00:00<00:00, 169363.95it/s]\n",
      "Processing text_left with extend: 100%|██████████| 2118/2118 [00:00<00:00, 842760.26it/s]\n",
      "Processing text_right with extend: 100%|██████████| 18841/18841 [00:00<00:00, 836304.08it/s]\n",
      "Building Vocabulary from a datapack.: 100%|██████████| 418401/418401 [00:00<00:00, 3460674.24it/s]\n",
      "Processing text_left with chain_transform of Tokenize => Lowercase => PuncRemoval: 100%|██████████| 2118/2118 [00:00<00:00, 10830.85it/s]\n",
      "Processing text_right with chain_transform of Tokenize => Lowercase => PuncRemoval: 100%|██████████| 18841/18841 [00:04<00:00, 4180.52it/s]\n",
      "Processing text_right with transform: 100%|██████████| 18841/18841 [00:00<00:00, 48605.24it/s]\n",
      "Processing text_left with transform: 100%|██████████| 2118/2118 [00:00<00:00, 134523.61it/s]\n",
      "Processing text_right with transform: 100%|██████████| 18841/18841 [00:00<00:00, 77349.04it/s]\n",
      "Processing length_left with len: 100%|██████████| 2118/2118 [00:00<00:00, 391621.23it/s]\n",
      "Processing length_right with len: 100%|██████████| 18841/18841 [00:00<00:00, 531842.50it/s]\n",
      "Processing text_left with chain_transform of Tokenize => Lowercase => PuncRemoval: 100%|██████████| 122/122 [00:00<00:00, 6614.60it/s]\n",
      "Processing text_right with chain_transform of Tokenize => Lowercase => PuncRemoval: 100%|██████████| 1115/1115 [00:00<00:00, 5589.66it/s]\n",
      "Processing text_right with transform: 100%|██████████| 1115/1115 [00:00<00:00, 66779.70it/s]\n",
      "Processing text_left with transform: 100%|██████████| 122/122 [00:00<00:00, 67480.56it/s]\n",
      "Processing text_right with transform: 100%|██████████| 1115/1115 [00:00<00:00, 67283.14it/s]\n",
      "Processing length_left with len: 100%|██████████| 122/122 [00:00<00:00, 105116.08it/s]\n",
      "Processing length_right with len: 100%|██████████| 1115/1115 [00:00<00:00, 317254.53it/s]\n",
      "Processing text_left with chain_transform of Tokenize => Lowercase => PuncRemoval: 100%|██████████| 237/237 [00:00<00:00, 5585.74it/s]\n",
      "Processing text_right with chain_transform of Tokenize => Lowercase => PuncRemoval: 100%|██████████| 2300/2300 [00:00<00:00, 4329.86it/s]\n",
      "Processing text_right with transform: 100%|██████████| 2300/2300 [00:00<00:00, 70847.65it/s]\n",
      "Processing text_left with transform: 100%|██████████| 237/237 [00:00<00:00, 121462.62it/s]\n",
      "Processing text_right with transform: 100%|██████████| 2300/2300 [00:00<00:00, 84147.30it/s]\n",
      "Processing length_left with len: 100%|██████████| 237/237 [00:00<00:00, 313382.74it/s]\n",
      "Processing length_right with len: 100%|██████████| 2300/2300 [00:00<00:00, 668067.81it/s]\n"
     ]
    }
   ],
   "source": [
    "train_pack_processed = preprocessor.fit_transform(train_pack_raw)\n",
    "dev_pack_processed = preprocessor.transform(dev_pack_raw)\n",
    "test_pack_processed = preprocessor.transform(test_pack_raw)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'filter_unit': <matchzoo.preprocessors.units.frequency_filter.FrequencyFilter at 0x7fa1cca00e90>,\n",
       " 'vocab_unit': <matchzoo.preprocessors.units.vocabulary.Vocabulary at 0x7fa1040e0a10>,\n",
       " 'vocab_size': 30058,\n",
       " 'embedding_input_dim': 30058}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "preprocessor.context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "glove_embedding = mz.datasets.embeddings.load_glove_embedding(dimension=300)\n",
    "term_index = preprocessor.context['vocab_unit'].state['term_index']\n",
    "embedding_matrix = glove_embedding.build_matrix(term_index)\n",
    "l2_norm = np.sqrt((embedding_matrix * embedding_matrix).sum(axis=1))\n",
    "embedding_matrix = embedding_matrix / l2_norm[:, np.newaxis]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainset = mz.dataloader.Dataset(\n",
    "    data_pack=train_pack_processed,\n",
    "    mode='pair',\n",
    "    num_dup=2,\n",
    "    num_neg=1\n",
    ")\n",
    "testset = mz.dataloader.Dataset(\n",
    "    data_pack=test_pack_processed\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "padding_callback = mz.models.aNMM.get_default_padding_callback()\n",
    "\n",
    "trainloader = mz.dataloader.DataLoader(\n",
    "    dataset=trainset,\n",
    "    batch_size=20,\n",
    "    stage='train',\n",
    "    resample=True,\n",
    "    sort=False,\n",
    "    shuffle=True,\n",
    "    callback=padding_callback,\n",
    ")\n",
    "testloader = mz.dataloader.DataLoader(\n",
    "    dataset=testset,\n",
    "    batch_size=20,\n",
    "    stage='dev',\n",
    "    sort=False,\n",
    "    shuffle=False,\n",
    "    callback=padding_callback,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "aNMM(\n",
      "  (embedding): Embedding(30058, 300)\n",
      "  (matching): Matching()\n",
      "  (hidden_layers): Sequential(\n",
      "    (0): Sequential(\n",
      "      (0): Linear(in_features=200, out_features=100, bias=True)\n",
      "      (1): ReLU()\n",
      "    )\n",
      "    (1): Sequential(\n",
      "      (0): Linear(in_features=100, out_features=1, bias=True)\n",
      "      (1): ReLU()\n",
      "    )\n",
      "  )\n",
      "  (q_attention): Linear(in_features=300, out_features=1, bias=False)\n",
      "  (dropout): Dropout(p=0.1)\n",
      "  (out): Linear(in_features=1, out_features=1, bias=True)\n",
      ") 9037903\n"
     ]
    }
   ],
   "source": [
    "model = mz.models.aNMM()\n",
    "\n",
    "model.params['task'] = ranking_task\n",
    "model.params['embedding'] = embedding_matrix\n",
    "\n",
    "model.params['dropout_rate'] = 0.1\n",
    "\n",
    "model.build()\n",
    "\n",
    "print(model, sum(p.numel() for p in model.parameters() if p.requires_grad))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9cf2558b03ac4601920281e324935653",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=102), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Iter-102 Loss-0.689]:\n",
      "  Validation: normalized_discounted_cumulative_gain@3(0.0): 0.4755 - normalized_discounted_cumulative_gain@5(0.0): 0.5304 - mean_average_precision(0.0): 0.4953\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f0425b2e65054556bf8949982f11dfc9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=102), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Iter-204 Loss-0.664]:\n",
      "  Validation: normalized_discounted_cumulative_gain@3(0.0): 0.4697 - normalized_discounted_cumulative_gain@5(0.0): 0.5263 - mean_average_precision(0.0): 0.4927\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cb63ec52db784fb093f943f4be9b99af",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=102), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Iter-306 Loss-0.575]:\n",
      "  Validation: normalized_discounted_cumulative_gain@3(0.0): 0.4958 - normalized_discounted_cumulative_gain@5(0.0): 0.5557 - mean_average_precision(0.0): 0.5199\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e226ab88b51b4887af6308341787aa4e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=102), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Iter-408 Loss-0.418]:\n",
      "  Validation: normalized_discounted_cumulative_gain@3(0.0): 0.5403 - normalized_discounted_cumulative_gain@5(0.0): 0.6118 - mean_average_precision(0.0): 0.5632\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3717f31733d54c8cb4486e6a286115d1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=102), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Iter-510 Loss-0.248]:\n",
      "  Validation: normalized_discounted_cumulative_gain@3(0.0): 0.582 - normalized_discounted_cumulative_gain@5(0.0): 0.6425 - mean_average_precision(0.0): 0.5923\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9be9ca16370a4e079caec2a649eb0aa8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=102), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Iter-612 Loss-0.135]:\n",
      "  Validation: normalized_discounted_cumulative_gain@3(0.0): 0.5938 - normalized_discounted_cumulative_gain@5(0.0): 0.6513 - mean_average_precision(0.0): 0.6101\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6556c3444da14ad4ab98e56e2b59e256",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=102), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Iter-714 Loss-0.087]:\n",
      "  Validation: normalized_discounted_cumulative_gain@3(0.0): 0.5984 - normalized_discounted_cumulative_gain@5(0.0): 0.6496 - mean_average_precision(0.0): 0.6069\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b31121aeb2e04d6886052c8b2713d59a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=102), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Iter-816 Loss-0.049]:\n",
      "  Validation: normalized_discounted_cumulative_gain@3(0.0): 0.5931 - normalized_discounted_cumulative_gain@5(0.0): 0.6466 - mean_average_precision(0.0): 0.6015\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "27df3eaa44e14247a8e3ba91b882c9a4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=102), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Iter-918 Loss-0.033]:\n",
      "  Validation: normalized_discounted_cumulative_gain@3(0.0): 0.5826 - normalized_discounted_cumulative_gain@5(0.0): 0.6466 - mean_average_precision(0.0): 0.6013\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fc8f2074ab7e436286a826b941947c3f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=102), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Iter-1020 Loss-0.023]:\n",
      "  Validation: normalized_discounted_cumulative_gain@3(0.0): 0.5803 - normalized_discounted_cumulative_gain@5(0.0): 0.6422 - mean_average_precision(0.0): 0.5978\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "29a31a2a0b194c55bb9379d69b5e1853",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=102), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Iter-1122 Loss-0.018]:\n",
      "  Validation: normalized_discounted_cumulative_gain@3(0.0): 0.5799 - normalized_discounted_cumulative_gain@5(0.0): 0.6406 - mean_average_precision(0.0): 0.5965\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2d5b0ab2d0e7452080757a0f08c3571d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=102), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Iter-1224 Loss-0.017]:\n",
      "  Validation: normalized_discounted_cumulative_gain@3(0.0): 0.566 - normalized_discounted_cumulative_gain@5(0.0): 0.6195 - mean_average_precision(0.0): 0.5762\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c106e36ccd03452db3113cc86179301e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=102), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Iter-1326 Loss-0.011]:\n",
      "  Validation: normalized_discounted_cumulative_gain@3(0.0): 0.5715 - normalized_discounted_cumulative_gain@5(0.0): 0.6304 - mean_average_precision(0.0): 0.5906\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1de668745f20473897eb2c2fa8eab0df",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=102), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Iter-1428 Loss-0.010]:\n",
      "  Validation: normalized_discounted_cumulative_gain@3(0.0): 0.5559 - normalized_discounted_cumulative_gain@5(0.0): 0.6198 - mean_average_precision(0.0): 0.5732\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6b4feb47ec874cff82f67e3873fe2ba1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=102), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Iter-1530 Loss-0.008]:\n",
      "  Validation: normalized_discounted_cumulative_gain@3(0.0): 0.557 - normalized_discounted_cumulative_gain@5(0.0): 0.6188 - mean_average_precision(0.0): 0.5757\n",
      "\n",
      "Cost time: 1121.3299095630646s\n"
     ]
    }
   ],
   "source": [
    "optimizer = torch.optim.Adam(model.parameters(), lr = 3e-4)\n",
    "\n",
    "trainer = mz.trainers.Trainer(\n",
    "    model=model,\n",
    "    optimizer=optimizer,\n",
    "    trainloader=trainloader,\n",
    "    validloader=testloader,\n",
    "    validate_interval=None,\n",
    "    epochs=15\n",
    ")\n",
    "\n",
    "trainer.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "match-zoo",
   "language": "python",
   "name": "match-zoo"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
